import io
import json
from typing import TypedDict, Optional

from PIL import Image
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langgraph.constants import END, START
from langgraph.graph import StateGraph

import llms
from tools import analysisQuestType, quest_xijielijie

llm = llms.getLLM()


class GraphState(TypedDict):
    """State of the graph."""
    articleId: Optional[str]  # 文章ID
    article: Optional[str]  # 文章
    questId: Optional[str]  # 题目ID
    quest: Optional[str]  # 题目内容
    type: Optional[str]  # 题目类型
    subType: Optional[str]  # 题目类型细分
    answer: Optional[str]  # 题目答案
    analysis: Optional[str]  # 原始解析
    finalAnalysis: Optional[str]  # 最终解析


def create_graph_state() -> GraphState:
    return {
        "quest": None,
        "article": None,
        "type": None,
        "answer": None,
        "analysis": None,
        "finalAnalysis": None,
        "questId": None,
        "articleId": None,
        "subType": None
    }


def setup_workflow() -> StateGraph:
    workflow = StateGraph(GraphState)
    workflow.add_node("analysisQuestType", analysisQuestType)
    workflow.add_node("quest_xijielijie", quest_xijielijie)  # 细节理解题

    # 题型路由
    def quest_type_router(state):
        if state.get('type') == "细节理解题":
            return "quest_xijielijie"
        return None

    workflow.add_conditional_edges(START,quest_type_router)

    # 定义图
    # workflow.set_entry_point("analysisQuestType")
    # workflow.add_edge("analysisQuestType", END)
    # workflow.set_finish_point("db_entry")
    return workflow


def main():
    workflow = setup_workflow()
    app = workflow.compile()

    im = Image.open(io.BytesIO(app.get_graph().draw_mermaid_png()))
    im.save("graph2.png")

    initial_state = create_graph_state()
    initial_state["articleId"] = "20280"
    initial_state["article"] = """"I spent my last vacation trekking in Nepal(尼泊尔), a small country in the Himalaya Mountains of Asia. On a trekking vacation, you walk for long distances(距离) through the mountains. There are very few roads in Nepal, so trekking is the best way to travel. It's best if you hire(雇佣) a guide. The guides know the best paths to take and they also help trekkers with their bags. Our guide even cooked our food for us. Every day, we trekked for about four hours before lunch, and about two hours in the afternoon. At the end of the day's trek, we sat around a campfire and talked about the things we saw. Then we went to sleep in large tents. Every morning, our guide brought breakfast to our tent. It was great. At the end of the trip, I didn't want to go home!


    """
    initial_state["quest"] = """
    ""2. Which statement is NOT true?
A．We slept in large tents.
B．The guide cooked food.
C．There are many roads in Nepal.
D．You can have breakfast in your tent."


    """
    initial_state["type"] = "细节理解题"
    initial_state["answer"] = "A"
    initial_state["analysis"] = """细节理解题。根据“白色、红色、黑色和黄色的裙子只要35美元。可知在米勒先生的店里能买到黄色的短裙。故选C"""

    # for event in app.stream(initial_state, stream_mode="values"):
    #     print(event)

    result = app.invoke(initial_state)
    json_result = json.dumps(result, indent=4, ensure_ascii=False)
    # print(json_result)
    print("【文章】")
    print(result["article"])
    print("【题干】")
    print(result["quest"])
    print("【答案】")
    print(result["answer"])
    print("【解析】")
    print(result["analysis"])
    print("【最终解析】")
    print(result["finalAnalysis"])

if __name__ == '__main__':
    main()
